import os
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()


def generate_sql(nl_query: str) -> str:
    """
    调用API，将自然语言查询转换SQL
    Args:
        nl_query (str): 自然语言查询字符串
    Returns:
        SQL字符串
    """
    table_schema = """
    表格名: sales
    字段:
    - 行 ID（整数类型）
    - 订单 ID（字符串类型）
    - 订单日期（日期时间类型，格式如 2018-04-27）
    - 发货日期（日期时间类型，格式如 2018-04-29）
    - 邮寄方式（字符串类型）
    - 客户 ID（字符串类型）
    - 客户名称（字符串类型）
    - 细分（字符串类型）
    - 城市（字符串类型）
    - 省/自治区（字符串类型）
    - 国家（字符串类型）
    - 地区（字符串类型）
    - 产品 ID（字符串类型）
    - 类别（字符串类型）
    - 子类别（字符串类型）
    - 产品名称（字符串类型，如 Fiskars 剪刀, 蓝色）
    - 销售额（浮点数类型，无单位标注）
    - 数量（整数类型）
    - 折扣（浮点数类型）
    - 利润（浮点数类型）
    """
    prompt = f"""
    先阅读表格结构信息，理解字段名及其类型。
    {table_schema}
    请将以下自然语言转换为查询上述表格的SQL语句，要求：
    1.只返回SQL，不包含任何解释
    2.严格使用表格中的字段名
    3.条件用WHERE，多条件用AND，分组用GROUP BY
    4.严格符合SQL语法规范
    
    自然语言查询: {nl_query}
    """

    try:
        client = OpenAI(
            # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
            api_key=os.getenv("DASHSCOPE_API_KEY"),
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )
        completion = client.chat.completions.create(
            model="qwen-plus",
            messages=[
                {"role": "user", "content": prompt}
            ]
        )
        return completion.choices[0].message.content.strip()
    except Exception as e:
        return f"生成失败: {str(e)}"


if __name__ == "__main__":
    test_query = "查一下2018-4-27杭州的订单信息"
    print("生成的SQL: ", generate_sql(test_query))
